package com.hospital.diagnosis.module.hanlder;

import cn.hutool.json.JSONUtil;
import com.hospital.common.base.ThreadAttributes;
import com.hospital.common.enums.RoleType;
import com.hospital.common.utils.LogUtils;
import com.hospital.common.utils.ThreadLocalUtils;
import com.hospital.diagnosis.module.base.PatientMessage;
import com.hospital.diagnosis.module.base.enums.TalkFlag;
import com.hospital.diagnosis.module.config.AiWebsocketProperties;
import com.hospital.diagnosis.module.exception.WebSocketException;
import com.hospital.diagnosis.module.exception.enums.WebSocketExpCode;
import com.hospital.diagnosis.module.utils.WebSocketSessionManager;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Objects;

import static com.hospital.common.base.interfaces.Const.*;

@Component
public class AiDiagnosisHandler extends TextWebSocketHandler {

    @Autowired
    private WebSocketSessionManager webSocketSessionManager;

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private AiWebsocketProperties aiWebsocketProperties;

    /**
     * 连接建立后触发,类比为@onOpen
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        ThreadAttributes threadAttributes = ThreadLocalUtils.get();
        checkAuth(threadAttributes);
        webSocketSessionManager.register(session.getId(), threadAttributes.getUserId());
        LogUtils.info("用户" + threadAttributes.getUserId() + "已连接");
    }

    /**
     * 处理发送的消息,类比为@onMessage
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        check(session);
        PatientMessage patientMessage = JSONUtil.toBean(message.getPayload(), PatientMessage.class);
        checkMessage(patientMessage);
        switch (Objects.requireNonNull(TalkFlag.getByValue(patientMessage.getTalkFlag()))) {
            case TALKING -> {
                String prompt = String.format("patientId为 %s，问题是：%s", patientMessage.getPatientId(), patientMessage.getMessage());
                chatClient.prompt()
                        .advisors(spec -> spec
                                .param(ChatMemory.CONVERSATION_ID, session.getId()))
                        .user(prompt)
                        .stream()
                        .content()
                        .toStream()
                        .forEach(part -> {
                            try {
                                session.sendMessage(new TextMessage(part));
                            } catch (IOException e) {
                                LogUtils.error("消息发送失败", e);
                            }
                        });
                session.sendMessage(new TextMessage(AI_MESSAGE_END_FLAG));
            }
            case FINISH -> {
                chatClient.prompt()
                        .advisors(spec -> spec
                                .param(ChatMemory.CONVERSATION_ID, session.getId()))
                        .user(aiWebsocketProperties.getConclusionPrompt())
                        .stream()
                        .content()
                        .toStream()
                        .forEach(part -> {
                            try {
                                session.sendMessage(new TextMessage(part));
                            } catch (IOException e) {
                                LogUtils.error("消息发送失败", e);
                            }
                        });
                session.sendMessage(new TextMessage(aiWebsocketProperties.getClosingPrompt()));
                session.sendMessage(new TextMessage(AI_MESSAGE_END_FLAG));
            }
        }
    }

    /**
     * 连接关闭后触发,类比为@onClose
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        LogUtils.info("会话:" + session.getId() + ",结束");
        webSocketSessionManager.removeSession(session.getId());
        chatClient.prompt()
                .advisors(spec -> spec
                        .param(ChatMemory.CONVERSATION_ID, session.getId()))
                .user(aiWebsocketProperties.getConclusionPrompt())
                .call();
    }

    /**
     * 处理会话产生的异常,类比为@onError
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws IOException {
        LogUtils.error("会话:" + session.getId() + ",异常", exception);
        session.sendMessage(new TextMessage(AI_MESSAGE_END_FLAG));
    }

    /**
     * 检查权限
     */
    private void checkAuth(ThreadAttributes threadAttributes){
        if(ObjectUtils.isEmpty(threadAttributes)){
            throw new WebSocketException(WebSocketExpCode.THREAD_ATTRIBUTES_NOT_EXIST);
        }
        if(!threadAttributes.getRoleType().equals(RoleType.PATIENT.value)){
            throw new WebSocketException(WebSocketExpCode.NOT_PATIENT);
        }
        if(ObjectUtils.isEmpty(threadAttributes.getUserId())){
            throw new WebSocketException(WebSocketExpCode.USER_NOT_EXIST);
        }
    }

    /**
     * 检查会话是否过期
     */
    private void check(WebSocketSession session) {
        if (webSocketSessionManager.expire(session.getId())) {
            LogUtils.error("用户会话已过期", new WebSocketException(WebSocketExpCode.SESSION_EXPIRED));
            webSocketSessionManager.removeSession(session.getId());
            try {
                session.close();
            } catch (IOException e) {
                throw new WebSocketException(WebSocketExpCode.SESSION_CLOSE_FAILED);
            }
            throw new WebSocketException(WebSocketExpCode.SESSION_EXPIRED);
        }
    }

    /**
     * 检查消息完整性
     */
    private void checkMessage(PatientMessage message) {
        if (ObjectUtils.isEmpty(message)) {
            throw new WebSocketException(WebSocketExpCode.MESSAGE_NOT_EXIST);
        }
        if (ObjectUtils.isEmpty(message.getTalkFlag())) {
            throw new WebSocketException(WebSocketExpCode.TALK_FLAG_NOT_EXIST);
        }
        if (ObjectUtils.isEmpty(TalkFlag.getByValue(message.getTalkFlag()))) {
            throw new WebSocketException(WebSocketExpCode.MESSAGE_FLAG_ERROR);
        }
        if (ObjectUtils.isEmpty(message.getPatientId())) {
            throw new WebSocketException(WebSocketExpCode.PATIENT_ID_NOT_EXIST);
        }
        if (ObjectUtils.isEmpty(message.getMessage())) {
            message.setMessage("");
        }
    }
}
